# Copyright (c) OpenMMLab. All rights reserved.
from .image import ImageClassifier

from ..builder import CLASSIFIERS

import torch

from globals import sets

# ------------------------------
# 2. 定义基于统计量对齐损失的函数
# ------------------------------
def alignment_loss_stats(mu_ft, Sigma_ft, mu_pre, Sigma_pre, Q):
    """
    计算统计量对齐损失：
      - 均值对齐： || mu_ft - mu_pre Q ||^2
      - 协方差对齐： || Sigma_ft - Q^T Sigma_pre Q ||_F^2
    """
    # loss_mean = torch.norm(mu_ft - mu_pre @ Q, p=2) ** 2/768
    loss_cov = (torch.norm(Sigma_ft - Q.t() @ Sigma_pre @ Q, p='fro') ** 2)/768/768
    # return loss_mean + loss_cov
    return loss_cov

@CLASSIFIERS.register_module()
class SimQ_ImageClassifier(ImageClassifier):
    def __init__(self,
                 backbone,
                 neck=None,
                 head=None,
                 pretrained=None,
                 train_cfg=None,
                 init_cfg=None,
                 mu_pre = None,
                 Sigma_pre = None,
                 Q = None):
        super().__init__(
                 backbone,
                 neck = neck,
                 head = head,
                 pretrained = pretrained,
                 train_cfg = train_cfg,
                 init_cfg = init_cfg)
        

        self.mu_pre = mu_pre
        self.Sigma_pre = Sigma_pre      
        self.Q = Q 


    def forward_train(self, img, gt_label=None, output_attentions=False, **kwargs):
        if self.augments is not None:
            img, gt_label = self.augments(img, gt_label)

        if output_attentions:
            x,attention_weights = self.extract_feat(img, output_attentions)
        else:
            x = self.extract_feat(img, output_attentions)
        if gt_label==None:
            if output_attentions:
                return x[-1], attention_weights
            return x[-1]
        
        losses = dict()
        loss = self.head.forward_train(x, gt_label)
        
        loss['loss'] = 1*loss['loss']
        if self.Q != None:
            feats = x[0]
            # 计算当前 batch 的统计量（特征均值和协方差）
            mu_ft = feats.mean(dim=0)  # (d,)
            centered = feats - mu_ft.unsqueeze(0)
            Sigma_ft = (centered.t() @ centered) / (feats.size(0) - 1)  # (d, d)
            
            # 基于统计量的对齐损失
            loss_align = alignment_loss_stats(mu_ft, Sigma_ft, self.mu_pre, self.Sigma_pre, self.Q)
            loss['align_loss'] = 1*loss_align
    


        losses.update(loss)

        return losses




